added inference model

app.py

@@ -1,7 +1,94 @@
-import gradio as gr
+import os 
+import torch
+import numpy as np
+import lightning as pl
+import gradio as gr 
+from PIL import Image
+from torchvision import transforms
+from data import  LitMNISTDataModule 
+from config import CONFIG
+from model import LitMNISTModel
+from timeit import default_timer as timer
+
+torch.set_float32_matmul_precision('medium')
+torch.cuda.amp.autocast(enabled=True)
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+torch.set_default_device( device= device  )
+
+
+pl.seed_everything(123, workers=True)
+
+TEST_TRANSFORMS = transforms.Compose([
+    # transforms.PILToTensor(),
+    # transforms.Resize((28, 28)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.1307,), (0.3081,))
+    ])
+
+## MNISTDataModule
+dm = LitMNISTDataModule(
+                data_dir=CONFIG['data'].get('dir_path','.'),
+                batch_size= CONFIG.get('batch_size'),
+                num_workers=CONFIG.get('num_workers'),
+                test_transform=TEST_TRANSFORMS,
+                train_transform=None
+    )
+dm.prepare_data()
+dm.setup('test')
+
+## MNISTModel
+# model =  LitMNISTModel()
+chkpoint_path = os.path.join(  os.path.dirname(__file__),'logs','chkpoints','epoch=13.ckpt' )
+model = LitMNISTModel.load_from_checkpoint(chkpoint_path)
+
+
+trainer  = pl.Trainer(
+                fast_dev_run=True,
+                precision=32,
+                enable_model_summary=False,
+                enable_progress_bar=False,
+             )
+
+# trainer.test(model,datamodule=dm)
+
+
+# for X,y in dm.test_dataloader():
+#     for i in range(X.shape[0]):
+#         plt.imsave(
+#             fname=os.path.join('numbers',f'img_{i}.png'),
+#             arr=np.clip(
+#                 torch.stack(
+#                     [X[i,...],X[i,...],X[i,...]],
+#                     dim=1
+#                 ).squeeze(0).permute(1,2,0).detach().cpu().contiguous().numpy(),0,1))
+#     break
+
+
+def predict_fn(img:Image):
+    start_time = timer()
+    try:    
+        img = np.array(img)        
+        img = TEST_TRANSFORMS(img)    
+        img = img.mean(dim=0).unsqueeze(0).unsqueeze(0).to(model.device)
+        y_preds = model.predict_step( img)
+        res = {f"Title: {y_preds['predict'][0]}": y_preds['prob'][0]}
+        pred_time = round(timer() - start_time, 5)
+        return(res, pred_time)
+    except Exception as e:
+        gr.Error("An error occured 💥!", duration=5)
+        return ({ f"Title ☠️": 0.0},0.0)
+
+
+
+
+gr.Interface(
+        fn=predict_fn,
+        inputs=gr.Image(type='pil'),
+        outputs=[
+            gr.Label(num_top_classes=1, label="Predictions"), # what are the outputs?
+            gr.Number(label="Prediction time (s)")
+        ],
+        examples=[ ['numbers/'+i] for i in os.listdir(os.path.join( os.path.dirname(__file__) ,'numbers'))]
+        ).launch(share=False,debug=False)
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()