update model

app.py

@@ -23,9 +23,13 @@ article = """
 <p>📦 Dataset <a href="https://github.com/eugenesiow/super-image-data">this GitHub repo</a></p>
 </div>
 """
+examples = [
+    ["LR_image.png"],
+    ["barbara.png"],
+]
 
 # load model
-print("Loading  SRCNN model...")
+# print("Loading  SRCNN model...")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 model = SRCNNModel().to(device)
@@ -33,7 +37,8 @@ model.load_state_dict(
     torch.load("SRCNNmodel_trained.pt", map_location=torch.device(device))
 )
 model.eval()
-print("SRCNN model loaded!")
+# print("SRCNN model loaded!")
+
 
 # def image_grid(imgs, rows, cols):
 #     '''
@@ -48,18 +53,17 @@ print("SRCNN model loaded!")
 #     for i, img in enumerate(imgs):
 #         grid.paste(img, box=(i%cols*w, i//cols*h))
 #     return grid
-examples = [
-    ["LR_image.png"],
-    ["barbara.png"],
-]
 
 
 def sepia(image):
     # gradio open image as np array
     image = Image.fromarray(image, mode="RGB")
-    out_final, image_bicubic, image = pred_SRCNN(
-        model=model, image=image, device=device
-    )
+
+    # prediction
+    with torch.no_grad():
+        out_final, image_bicubic, image = pred_SRCNN(
+            model=model, image=image, device=device
+        )
     # grid = image_grid([out_final,image_bicubic],1,2)
     return out_final, image_bicubic
 

model.py

@@ -59,22 +59,22 @@ def pred_SRCNN(model,image,device,scale_factor=2):
     return out_final,image_bicubic,image
 
 
-def main():
-    print("Loading  SRCNN model...")
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# def main():
+#     print("Loading  SRCNN model...")
+#     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-    model = SRCNNModel().to(device)
-    model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
-    model.eval()
-    print("SRCNN model loaded!")
+#     model = SRCNNModel().to(device)
+#     model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
+#     model.eval()
+#     print("SRCNN model loaded!")
 
-    image_path = "LR_image.png"
+#     image_path = "LR_image.png"
 
-    out_final,image_bicubic,image = pred_SRCNN(model=model,image_path=image_path,device=device)
-    image.show()
-    out_final.show()
-    image_bicubic.show()
+#     out_final,image_bicubic,image = pred_SRCNN(model=model,image_path=image_path,device=device)
+#     image.show()
+#     out_final.show()
+#     image_bicubic.show()
 
 
-if __name__=="__main__":
-    main()
+# if __name__=="__main__":
+#     main()