Update app.py

app.py

@@ -22,7 +22,7 @@ with open(config_path, 'r') as data:
     config = yaml.safe_load(data)
 system = System(config)
 checkpoint_path = 'weights/parseq.ckpt'
-checkpoint = torch.load(checkpoint_path, map_location = 'cpu')
+checkpoint = torch.load(checkpoint_path, map_location = 'cuda')
 system.load_state_dict(checkpoint['state_dict'])
 system.to(device)
 
@@ -30,13 +30,20 @@ def predict(image):
     if isinstance(image, str):
         image = cv2.imread(image)
     _, img_wapred, _, _ = detect_lp(wpod_net, image, 0.5)
-    img =  (img_wapred[0] * 255).astype(np.uint8)
-    img = Image.fromarray(img).convert("RGB")
-    image = trans(img).unsqueeze(0)
-    with torch.no_grad():
-        pred = system(image).softmax(-1)
-    generated_text, _ = system.tokenizer.decode(pred)
-    return generated_text[0]
+    if len(img_warped) == 0:
+        return "Can not detect license plate from image"
+    else:
+        system.eval()
+        pred_labels = []
+        for i in range(len(img_warped)):
+            img =  (img_wapred[i] * 255).astype(np.uint8)
+            img = Image.fromarray(img).convert("RGB")
+            image = trans(img).unsqueeze(0)
+            with torch.no_grad():
+                pred = system(image).softmax(-1)
+            generated_text, _ = system.tokenizer.decode(pred)
+            pred_labels.append(generated_text[0])
+        return pred_labels
 
 interface = gr.Interface(
     fn = predict,