Update app.py

app.py

@@ -13,7 +13,8 @@ import torch.nn.functional as F
 import numpy as np
 import matplotlib.pyplot as plt
 import random
-from skimage.transform import resize
+import zipfile
+# from skimage.transform import resize
 from torchvision import transforms, models
 
 
@@ -302,23 +303,19 @@ def create_retrieval_figure(res):
     ax_query.set_title('Top 10 most similar scarabs', fontsize=40)
     names = ""
     for i, image in zip(range(len(res)), res):
-        current_image_path = image
+        current_image_path = image.split("/")[3]+"/"+image.split("/")[4]
         if i==0: continue
         if i < 11:
-            image = cv2.imread(current_image_path)
+            archive = zipfile.ZipFile('dataset.zip', 'r')
+            imgfile = archive.read(current_image_path)
+            image = cv2.imdecode(np.frombuffer(imgfile, np.uint8), 1)
             # image_resized = cv2.resize(image, (224, 224), interpolation=cv2.INTER_LINEAR)
             ax = fig.add_subplot(rows, cols, i)
             plt.axis('off')
             plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-            item_uuid = current_image_path.split("/")[4].split("_photoUUID")[0].split("itemUUID_")[1]
+            item_uuid = current_image_path.split("/")[1].split("_photoUUID")[0].split("itemUUID_")[1]
             ax.set_title('Top {}'.format(i), fontsize=40)
             names = names + "Top " + str(i) + " item UUID is " + item_uuid + "\n"
-    # img_buf = io.BytesIO()
-    # plt.savefig(img_buf, format='png')
-    # im_fig = Image.open(img_buf)
-    # img_buf.close()
-    # return im_fig
-
     return fig, names
 
 def knn_calc(image_name, query_feature, features):
@@ -328,38 +325,29 @@ def knn_calc(image_name, query_feature, features):
     dist = -dist.item()
     return dist
 
-# device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-device = 'cpu'
-
-experiment = "experiment_0"
-checkpoint_path = os.path.join("../shapes_classification/checkpoints/"
-                               "50_50_pretrained_resnet101_experiment_0_train_images_with_drawings_batch_8_10:29:06/" +
-                              "experiment_0_last_auto_model.pth.tar")
 checkpoint_path = "multi_label.pth.tar"
 
 resnet = models.resnet101(pretrained=True)
 num_ftrs = resnet.fc.in_features
 resnet.fc = nn.Linear(num_ftrs, 13)
-model = Resnet_with_skip(resnet).to(device)
+model = Resnet_with_skip(resnet)
 checkpoint = torch.load(checkpoint_path, map_location="cpu")
 model.load_state_dict(checkpoint)
 embedding_model_test = torch.nn.Sequential(*(list(model.children())[:-1]))
-embedding_model_test.to(device)
 
 periods_model = models.resnet101(pretrained=True)
 periods_model.fc = nn.Linear(num_ftrs, 5)
 periods_checkpoint = torch.load("periods.pth.tar", map_location="cpu")
 periods_model.load_state_dict(periods_checkpoint)
-periods_model.to(device)
 
-data_dir = "../cssl_dataset/all_image_base/1/"
-query_images_paths = []
-for path in os.listdir(data_dir):
-    query_images_paths.append(os.path.join(data_dir, path))
+with open('query_images_paths.pkl', 'rb') as fp:
+    query_images_paths = pickle.load(fp)
 
 with open('features.pkl', 'rb') as fp:
     features = pickle.load(fp)
 
+
+
 model.eval()
 transform = transforms.Compose([
             transforms.Resize((224, 224)),
@@ -373,21 +361,21 @@ invTrans = transforms.Compose([transforms.Normalize(mean=[0., 0., 0.],
                                                     std=[1., 1., 1.]),
                                ])
 
-labels = sorted(os.listdir("../cssl_dataset/shape_multi_label/photos"))
+labels = ['ankh', 'anthropomorphic', 'bands', 'beetle', 'bird', 'circles', 'cross', 'duck', 'head', 'ibex', 'lion', 'sa', 'snake']
+
 periods_labels = ["MB1", "MB2", "LB", "Iron1", 'Iron2']
 periods_model.eval()
 
 def predict(inp):
     image_tensor = transform(inp)
-    image_tensor = image_tensor.to(device)
     with torch.no_grad():
         classification, reconstruction = model(image_tensor.unsqueeze(0))
         periods_classification = periods_model(image_tensor.unsqueeze(0))
         recon_tensor = reconstruction[0].repeat(3, 1, 1)
         recon_tensor = invTrans(kornia.enhance.invert(recon_tensor))
-        plot_recon = recon_tensor.to("cpu").permute(1, 2, 0).detach().numpy()
+        plot_recon = recon_tensor.permute(1, 2, 0).detach().numpy()
         w, h = inp.size
-        plot_recon = resize(plot_recon, (h, w))
+        # plot_recon = resize(plot_recon, (h, w))
         m = nn.Sigmoid()
         y = m(classification)
         preds = []
@@ -397,7 +385,6 @@ def predict(inp):
                     preds.append(1)
                 else:
                     preds.append(0)
-        # prediction = torch.tensor(preds).to(device)
         confidences = {}
         true_labels = ""
         for i in range(len(labels)):
@@ -406,11 +393,11 @@ def predict(inp):
                     true_labels = true_labels + labels[i]
                 else:
                     true_labels = true_labels + "&" + labels[i]
-        confidences[true_labels] = torch.tensor(1.0).to(device)
+        confidences[true_labels] = torch.tensor(1.0)
 
         periods_prediction = torch.nn.functional.softmax(periods_classification[0], dim=0)
         periods_confidences = {periods_labels[i]: periods_prediction[i] for i in range(len(periods_labels))}
-        feature = embedding_model_test(image_tensor.unsqueeze(0)).to(device)
+        feature = embedding_model_test(image_tensor.unsqueeze(0))
         dists = dict()
         with torch.no_grad():
             for i, image_name in enumerate(query_images_paths):
@@ -418,9 +405,14 @@ def predict(inp):
                 dists[image_name] = dist
         res = dict(sorted(dists.items(), key=itemgetter(1)))
     fig, names = create_retrieval_figure(res)
-    return fig, names, plot_recon, confidences, periods_confidences
+    return confidences, periods_confidences, plot_recon, fig, names
 
 
 gr.Interface(fn=predict,
              inputs=gr.Image(type="pil"),
-             outputs=['plot', 'text', "image", gr.Label(num_top_classes=1), gr.Label(num_top_classes=1)], ).launch(share=True)
+             title="ArcAid: Analysis of Archaeological Artifacts using Drawings",
+             description="Easily classify artifacs, retrieve similar ones and generate drawings. "
+                         "https://arxiv.org/abs/2211.09480.",
+             examples=['anth.jpg', 'beetle_snakes.jpg', 'bird.jpg', 'cross.jpg', 'ibex.jpg',
+                       'lion.jpg', 'lion2.jpg', 'sa.jpg'],
+             outputs=[gr.Label(num_top_classes=1), gr.Label(num_top_classes=1), "image", 'plot', 'text', "image"], ).launch(share=True)