workflow stuff

checkpoints/{ham10k_checkpoint_mobile_0.82_epoch24.pt → state_dict_timm_effnet_b3_e6_val_f1=0.75.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7c671f87198d4c5bb3a442ca9bf810dad45f755f6f7045b0b7b5186df9f767db
-size 50851683
+oid sha256:7e3ba0141dacacd6a79a6073f5e6bbb6daa3f132944635b54adb7c94f2761e99
+size 43381701

runSDSdemo.py

@@ -8,18 +8,36 @@ import torchvision.transforms as transforms
 from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
 from pytorch_grad_cam.utils.image import show_cam_on_image
 import gradio as gr
+import timm
 
 # model setup
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-classes = [ 'actinic keratoses', 'basal cell carcinoma', 'benign keratosis-like lesions', 
-           'dermatofibroma','melanoma', 'melanocytic nevi', 'vascular lesions']
-model = torchvision.models.mobilenet_v3_large(pretrained = False) # This is a very well known network but it is designed for 1000 classes and not just cats and dogs this is why we need the next line
-model.classifier[3] = nn.Linear(1280, 7)
+classes = ["Preparation",
+        "CalotTriangleDissection",
+        "ClippingCutting",
+        "GallbladderDissection",
+        "GallbladderPackaging",
+        "CleaningCoagulation",
+        "GallbladderRetraction"]
+
+
+model = timm.create_model('efficientnet_b3')# This is a very well known network but it is designed for 1000 classes and not just cats and dogs this is why we need the next line
+model.classifier = nn.Linear(1536, 7)
 #state_dict_trained = torch.hub.load_state_dict_from_url("https://github.com/tobiascz/demotime/raw/main/checkpoints/ham10k_checkpoint_mobile_0.82_epoch24.pt", model_dir=".", map_location = device)
 import os
 print(os.getcwd())
-state_dict_trained = torch.load('checkpoints/ham10k_checkpoint_mobile_0.82_epoch24.pt', map_location=torch.device('cpu'))
-model.load_state_dict(state_dict_trained["model_state_dict"]) ## Here we load the trained weights (state_dict) in our model 
+state_dict_trained = torch.load('checkpoints/state_dict_timm_effnet_b3_e6_val_f1=0.75.pt', map_location=torch.device(device))
+sd = model.state_dict()
+
+print(state_dict_trained.keys())
+
+for k,v in sd.items():
+      if not "classifier" in k:
+        sd[k] = state_dict_trained[f'model.model.{k}']
+sd['classifier.weight'] = state_dict_trained['model.fc_phase.weight']
+sd['classifier.bias'] = state_dict_trained['model.fc_phase.bias']
+      
+model.load_state_dict(sd) ## Here we load the trained weights (state_dict) in our model 
 model.eval() # This
 
 # image pre-processing
@@ -76,21 +94,18 @@ def predict(input_img):
   # class with hightest probability
   pred = torch.argmax(outputs, dim=1).cpu().numpy()
   # diagnostic suggestions
-  if pred == 1 or pred == 4:
-    suggestion = "CHECK WITH YOUR MD!"
-  else:
-    suggestion = "Nothing to be worried about."
+
   # grad_cam image
   target_layers = model.features[-1]
   output_img = image_grad_cam(model,leasion_tensor,input_float_np,target_layers)
   # return label dict and suggestion
-  return {classes[i]: float(pred_softmax[i]) for i in range(len(classes))}, suggestion, output_img
+  return {classes[i]: float(pred_softmax[i]) for i in range(len(classes))}, output_img
 
 # start gradio application
 gr.Interface(
         fn=predict, 
         inputs=gr.inputs.Image(), 
-        outputs=[gr.outputs.Label(label="Predict Result"), gr.outputs.Textbox(type="str", label="Recommendation"), gr.outputs.Image(label="GradCAM")],
-        examples=[['sample1.png'],['sample2.jpg'],['sample3.jpg']],
-        title="Skin Lesion Classifier"
+        outputs=[gr.outputs.Label(label="Predict Result"), gr.outputs.Image(label="GradCAM")],
+        examples=[['images/video01_000014_prep.png'],['images/video01_001403.png'],['images/video01_001528_pack.png']],
+        title="Surgical Workflow Classifier"
       ).launch()