Update app.py

app.py

@@ -1,170 +1,56 @@
 import os
-import time
-import argparse
-import numpy as np
-import random
 import pandas as pd
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision
-import sys
-import json
-from collections import defaultdict
-import math
 import gradio as gr
-
 from model import DistMult
-
-from tqdm import tqdm
-from utils import collate_list, detach_and_clone, move_to
 from PIL import Image
 from torchvision import transforms
+import json
+from tqdm import tqdm
 
+# Default image tensor normalization
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_MEAN = [0.485, 0.456, 0.406]
 _DEFAULT_IMAGE_TENSOR_NORMALIZATION_STD = [0.229, 0.224, 0.225]
 
-def evaluate(img):
-    num_ent_id = len(id2entity)
-    model = DistMult(args, num_ent_id, target_list, args.device)
-    model.to(args.device)
-    model.eval()
-    
-    torch.set_grad_enabled(False)
+# Load necessary data and initialize the model
+entity2id = json.load(open('entity2id_subtree.json', 'r'))
+id2entity = {v: k for k, v in entity2id.items()}
+datacsv = pd.read_csv('dataset_subtree.csv', low_memory=False)
+num_ent_id = len(entity2id)
+target_list = generate_target_list(datacsv, entity2id)  # Assuming this function is defined elsewhere
 
-    overall_id_to_name = json.load(open('overall_id_to_name.json'))
+# Initialize your model here
+model = DistMult(args, num_ent_id, target_list, torch.device('cpu'))  # Update arguments as necessary
+model.eval()
 
-    img = Image.open(args.img_path).convert('RGB')
-            
-    transform_steps = transforms.Compose([transforms.Resize((448, 448)), transforms.ToTensor(), transforms.Normalize(_DEFAULT_IMAGE_TENSOR_NORMALIZATION_MEAN, _DEFAULT_IMAGE_TENSOR_NORMALIZATION_STD)])
+# Define your evaluation function
+def evaluate(img):
+    transform_steps = transforms.Compose([
+        transforms.Resize((448, 448)),
+        transforms.ToTensor(),
+        transforms.Normalize(_DEFAULT_IMAGE_TENSOR_NORMALIZATION_MEAN, _DEFAULT_IMAGE_TENSOR_NORMALIZATION_STD)
+    ])
     h = transform_steps(img)
     r = torch.tensor([3])
 
-    h = move_to(h, args.device).unsqueeze(0)
-    r = move_to(r, args.device).unsqueeze(0)
+    # Assuming `move_to` is a function to move tensors to the desired device
+    h = h.unsqueeze(0)
+    r = r.unsqueeze(0)
 
     outputs = model.forward_ce(h, r, triple_type=('image', 'id'))
-
-    y_pred = detach_and_clone(outputs.cpu())
-    y_pred = y_pred.argmax(-1)
-
+    y_pred = outputs.argmax(-1).cpu()
     pred_label = target_list[y_pred].item()
     species_label = overall_id_to_name[str(id2entity[pred_label])]
-    print('species label = {}'.format(species_label))
-
-    # predict multi-level classification
-    
-# def get_classification(img):
-
-#     image_tensor = transform_image(img)
-#     ort_inputs = {input_name: to_numpy(image_tensor)}
-#     x = ort_session.run(None, ort_inputs)
-#     predictions = torch.topk(torch.from_numpy(x[0]), k=5).indices.squeeze(0).tolist()
-
-#     result = {}
-#     for i in predictions:
-#         label = label_map[str(i)]
-#         prob = x[0][0, i].item()
-#         result[label] = prob
-#     return result
 
-
-
-# iface.launch()
-
-    
     return species_label
 
-def _get_id(dict, key):
-    id = dict.get(key, None)
-    if id is None:
-        id = len(dict)
-        dict[key] = id
-    return id
-
-def generate_target_list(data, entity2id):
-    sub = data.loc[(data["datatype_h"] == "image") & (data["datatype_t"] == "id"), ['t']]
-    sub = list(sub['t'])
-    categories = []
-    for item in tqdm(sub):
-        if entity2id[str(int(float(item)))] not in categories:
-            categories.append(entity2id[str(int(float(item)))])
-    # print('categories = {}'.format(categories))
-    # print("No. of target categories = {}".format(len(categories)))
-    return torch.tensor(categories, dtype=torch.long).unsqueeze(-1)
-
-
-
-if __name__=='__main__':
-    parser = argparse.ArgumentParser()
-    # parser.add_argument('--data-dir', type=str, default='data/iwildcam_v2.0/')
-    # parser.add_argument('--img-path', type=str, required=True, help='path to species image to be classified')
-    parser.add_argument('--seed', type=int, default=813765)
-    parser.add_argument('--ckpt-path', type=str, default=None, help='path to ckpt for restarting expt')
-    parser.add_argument('--debug', action='store_true')
-    parser.add_argument('--no-cuda', action='store_true')
-    parser.add_argument('--batch_size', type=int, default=16)
-
-    parser.add_argument('--embedding-dim', type=int, default=512)
-    parser.add_argument('--location_input_dim', type=int, default=2)
-    parser.add_argument('--time_input_dim', type=int, default=1)
-    parser.add_argument('--mlp_location_numlayer', type=int, default=3)
-    parser.add_argument('--mlp_time_numlayer', type=int, default=3)
-
-    parser.add_argument('--img-embed-model', choices=['resnet18', 'resnet50'], default='resnet50')
-    parser.add_argument('--use-data-subset', action='store_true')
-    parser.add_argument('--subset-size', type=int, default=10)
-
-    args = parser.parse_args()
-
-    print('args = {}'.format(args))
-    args.device = torch.device('cuda') if not args.no_cuda and torch.cuda.is_available() else torch.device('cpu')
-
-    # Set random seed
-    torch.manual_seed(args.seed)
-    np.random.seed(args.seed)
-    random.seed(args.seed)
-
-    datacsv = pd.read_csv('dataset_subtree.csv', low_memory=False)
-
-    entity_id_file = 'entity2id_subtree.json'
-
-    if not os.path.exists(entity_id_file):
-        entity2id = {} # each of triple types have their own entity2id
-        
-        for i in tqdm(range(datacsv.shape[0])):
-            if datacsv.iloc[i,1] == "id":
-                _get_id(entity2id, str(int(float(datacsv.iloc[i,0]))))
-
-            if datacsv.iloc[i,-2] == "id":
-                _get_id(entity2id, str(int(float(datacsv.iloc[i,-3]))))
-        json.dump(entity2id, open(entity_id_file, 'w'))
-    else:
-        entity2id = json.load(open(entity_id_file, 'r'))
-    
-    id2entity = {v:k for k,v in entity2id.items()}
-
-    num_ent_id = len(entity2id)
-
-    # print('len(entity2id) = {}'.format(len(entity2id)))
-    
-    target_list = generate_target_list(datacsv, entity2id)
-
-    # restore from ckpt
-    if args.ckpt_path:
-        ckpt = torch.load(args.ckpt_path, map_location=args.device)
-        model.load_state_dict(ckpt['model'], strict=False)
-        print('ckpt loaded...')
-
-    species_model = gr.Interface(
-        evaluate,
-        gr.inputs.Image(shape=(200, 200)),
-        outputs="label",
-        title = 'Species Classification',
-        description = 'Species Classification',
-        article = 'Species Classification',
-        name = 'Species Classification',
-    )
-    species_model.launch(server_port=7897)
-
-    # evaluate(model, id2entity, target_list, args)
+# Gradio interface
+species_model = gr.Interface(
+    evaluate,
+    gr.inputs.Image(shape=(200, 200)),
+    outputs="label",
+    title='Species Classification',
+    description='Species Classification',
+    article='Species Classification'
+)
+species_model.launch()