Upload app.py

app.py

@@ -20,10 +20,10 @@
 
 # Libraries
 import torch
-import pickle
 from utils import *
 import gradio as gr
 from numpy import array
+from darknet import Darknet
 from torch.autograd import Variable
 from torch.cuda import is_available as check_cuda
 from PIL.ImageOps import grayscale
@@ -37,7 +37,9 @@ nms_thresh = 0.30
 run_cuda = False
 
 # CFG Files
+cfg     = 'cfg/yolov3-openimages.cfg'
 clsnames= 'cfg/openimages.names'
+weights = 'cfg/yolov3-openimages.weights'
 
 # Load classes
 classes = load_classes(clsnames)
@@ -45,8 +47,10 @@ num_classes = len(classes)
 
 # Set up the neural network
 print('Load Network')
-with open("models/darknet.pkl", 'rb') as f:
-    model = pickle.load(f)
+model = Darknet(cfg)
+
+print('Load Weights')
+model.load_weights(weights)
 
 print('Successfully loaded Network')
 
@@ -63,6 +67,9 @@ inp_dim = int(model.net_info["height"])
 if CUDA:
     model.cuda()
 
+# Set the model in evaluation mode
+model.eval()
+
 def get_detections(x):
     c1 = [int(y) for y in x[1:3]]
     c2 = [int(y) for y in x[3:5]]
@@ -248,7 +255,7 @@ def predict(img):
         emotions = {learn_emotion_labels[i]: float(probs_emotion[i]) for i in range(len(learn_emotion_labels))}
         sentiments = {learn_sentiment_labels[i]: float(probs_sentiment[i]) for i in range(len(learn_sentiment_labels))}
 
-        output = [img, emotions, sentiments, None, None, None, None, None, None]
+        output = [img.resize((12, 12)), emotions, sentiments, None, None, None, None, None, None]
     
     else: # Max 3 for now
         for face in faces[:3]:
@@ -262,7 +269,7 @@ def predict(img):
             emotions = {learn_emotion_labels[i]: float(probs_emotion[i]) for i in range(len(learn_emotion_labels))}
             sentiments = {learn_sentiment_labels[i]: float(probs_sentiment[i]) for i in range(len(learn_sentiment_labels))}
 
-            output.append(img)
+            output.append(img.resize((12, 12)))
             output.append(emotions)
             output.append(sentiments)