Upload batched_inference.py

batched_inference.py

@@ -0,0 +1,193 @@
+import csv
+import torch.multiprocessing as multiprocessing
+import pandas as pd
+import numpy as np
+import torchvision.transforms as transforms
+from torch import autocast
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+import torch
+from torchvision.transforms import InterpolationMode
+from tqdm import tqdm
+import random
+import json
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+torch.autograd.set_detect_anomaly(False)
+
+torch.autograd.profiler.emit_nvtx(enabled=False)
+torch.autograd.profiler.profile(enabled=False)
+torch.backends.cudnn.benchmark = True
+
+class ImageDataset(Dataset):
+    def __init__(self, csv_file, train, base_path):
+
+        self.csv_file = csv_file
+        self.train = train
+        self.all_image_names = self.csv_file[:]['md5'].apply(str)
+        self.all_image_ext = self.csv_file[:]['file_ext'].apply(str)
+        self.train_size = len(self.csv_file)
+        self.base_path = base_path
+        if self.train == True:
+            print(f"Number of training images: {self.train_size}")
+            self.thin_transform = transforms.Compose([
+                transforms.Resize(224, interpolation=InterpolationMode.BICUBIC),
+                transforms.CenterCrop(224),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[
+                    0.48145466,
+                    0.4578275,
+                    0.40821073
+                ], std=[
+                    0.26862954,
+                    0.26130258,
+                    0.27577711
+                ])  # Normalize image
+
+            ])
+            self.normal_transform = transforms.Compose([
+                transforms.Resize((224, 224), interpolation=InterpolationMode.BICUBIC),
+                transforms.ToTensor(),
+                transforms.Normalize(mean=[
+                    0.48145466,
+                    0.4578275,
+                    0.40821073
+                ], std=[
+                    0.26862954,
+                    0.26130258,
+                    0.27577711
+                ])  # Normalize image
+
+            ])
+
+    def __len__(self):
+        return len(self.all_image_names)
+
+    def __getitem__(self, index):
+        image = Image.open(self.base_path+"/"+str(self.all_image_names[index])+str(self.all_image_ext[index])).convert("RGB")
+        ratio = image.height/image.width
+        if ratio > 2.0 or ratio < 0.5:
+            image = self.thin_transform(image)
+        else:
+            image = self.normal_transform(image)
+
+
+        return {
+            'image': image,
+            "image_name": self.all_image_names[index]
+        }
+
+
+def prepare_model():
+    model = torch.load("path/to/your/model.pth").to("cuda")
+    model.to(memory_format=torch.channels_last)
+    model = model.eval()
+    return model
+
+
+def train(tagging_is_running, model, dataloader, train_data, output_queue):
+    print('Begin tagging')
+    model.eval()
+    counter = 0
+
+    with torch.no_grad():
+        for i, data in tqdm(enumerate(dataloader), total=int(len(train_data) / dataloader.batch_size)):
+
+            data, image_names = data['image'].to("cuda"), data["image_name"]
+            with autocast(device_type='cuda', dtype=torch.bfloat16):
+                outputs = model(data)
+
+                probabilities = torch.nn.functional.sigmoid(outputs)
+                output_queue.put((probabilities.to("cpu"), image_names))
+
+            counter += 1
+    _ = tagging_is_running.get()
+    print("Tagging finished!")
+
+
+def tag_writer(tagging_is_running, output_queue, output_file_name):
+    with open("tags.json", "r") as file:
+        tags = json.load(file)
+    allowed_tags = sorted(tags)
+    del tags
+    allowed_tags.extend(["placeholder0", "placeholder1", "placeholder2"])
+    tag_count = len(allowed_tags)
+    assert tag_count == 7704, f"The length of loss scaling factor is not correct. Correct: 7704, current: {tag_count}"
+
+    with open(output_file_name, "w") as output_csv:
+        writer = csv.writer(output_csv)
+        writer.writerow(["image_name", "tags", "tag_probs"])
+        while not (tagging_is_running.qsize()>0 and output_queue.qsize()>0):
+            tag_probabilities, image_names = output_queue.get()
+            tag_probabilities = tag_probabilities.tolist()
+
+            for per_image_tag_probabilities,image_name  in zip(tag_probabilities, image_names, strict=True):
+                this_image_tags = []
+                this_image_tag_probabilities = []
+                for index, per_tag_probability in enumerate(per_image_tag_probabilities):
+                    if per_tag_probability > 0.3:
+                        tag = allowed_tags[index]
+                        if "placeholder" not in tag:
+                            this_image_tags.append(tag)
+                            this_image_tag_probabilities.append(str(int(round(per_tag_probability, 3) * 1000)))
+                image_row = [image_name," ".join(this_image_tags)," ".join(this_image_tag_probabilities)]
+                writer.writerow(image_row)
+
+
+
+
+
+
+def set_seed(seed: int = 42) -> None:
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    # When running on the CuDNN backend, two further options must be set
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    # Set a fixed value for the hash seed
+    print(f"Random seed set as {seed}")
+
+
+if __name__ == "__main__":
+    steps = 0
+    output_file_name = "your_file.csv"
+    set_seed()
+    multiprocessing.set_start_method('spawn')
+    output_queue = multiprocessing.Queue()
+    tagging_is_running = multiprocessing.Queue(maxsize=5)
+    tagging_is_running.put("Running!")
+
+    # initialize the computation device
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        raise RuntimeError("CUDA is not available!")
+
+    model = prepare_model().to("cuda")
+    batch_size = 128
+
+
+    # read the training csv file
+    train_csv = pd.read_csv('/path/to/a/list/of/files/and/their/extensions.csv')
+    # train dataset
+    train_data = ImageDataset(
+        train_csv, train=True
+    )
+
+    train_loader = DataLoader(
+        train_data,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=6,
+        pin_memory=True
+    )
+    process_writer = multiprocessing.Process(target=tag_writer, args=(tagging_is_running, output_queue, output_file_name))
+    process_writer.start()
+    process_tagger = multiprocessing.Process(target=train, args=(tagging_is_running, model, train_loader, train_data, output_queue,))
+    process_tagger.start()
+    process_writer.join()
+    process_tagger.join()