Update app.py

app.py

@@ -38,112 +38,7 @@ from utils import (
     get_loaders
 )
 from loss import YoloLoss
-
-
-# custom functions for yolo
-# loss function for yolov3
-loss_fn = YoloLoss()
-
-def model_criterion(out, y,anchors):
-    loss = (  loss_fn(out[0], y[0], anchors[0])
-            + loss_fn(out[1], y[1], anchors[1])
-            + loss_fn(out[2], y[2], anchors[2])
-            )
-    return loss
-
-
-# accuracy function for yolov3
-def accuracy_fn(y, out, threshold,correct_class, correct_obj,correct_noobj, tot_class_preds,tot_obj, tot_noobj):
-
-    for i in range(3):
-
-      obj = y[i][..., 0] == 1 # in paper this is Iobj_i
-      noobj = y[i][..., 0] == 0  # in paper this is Iobj_i
-
-      correct_class += torch.sum(
-          torch.argmax(out[i][..., 5:][obj], dim=-1) == y[i][..., 5][obj]
-      )
-      tot_class_preds += torch.sum(obj)
-
-      obj_preds = torch.sigmoid(out[i][..., 0]) > threshold
-      correct_obj += torch.sum(obj_preds[obj] == y[i][..., 0][obj])
-      tot_obj += torch.sum(obj)
-      correct_noobj += torch.sum(obj_preds[noobj] == y[i][..., 0][noobj])
-      tot_noobj += torch.sum(noobj)
-
-    return((correct_class/(tot_class_preds+1e-16))*100,
-         (correct_noobj/(tot_noobj+1e-16))*100,
-         (correct_obj/(tot_obj+1e-16))*100)
-
-# pytorch lightning
-class LitYolo(LightningModule):
-    def __init__(self, num_classes=config.NUM_CLASSES, lr=1E-3,weight_decay=config.WEIGHT_DECAY,threshold=config.CONF_THRESHOLD):
-        super().__init__()
-
-        self.save_hyperparameters()
-        self.model = YOLOv3(num_classes=self.hparams.num_classes)
-        self.criterion = model_criterion
-        self.accuracy_fn = accuracy_fn
-        self.scaled_anchors = (torch.tensor(config.ANCHORS) * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2))
-        self.tot_class_preds, self.correct_class = 0, 0
-        self.tot_noobj, self.correct_noobj = 0, 0
-        self.tot_obj, self.correct_obj = 0, 0
-
-    def forward(self, x):
-        out = self.model(x)
-        return out
-
-    def training_step(self, batch, batch_idx):
-        x, y = batch
-        out  = self(x)
-        loss = self.criterion(out,y,self.scaled_anchors)
-        acc  = self.accuracy_fn(y,out,self.hparams.threshold,self.correct_class,
-                                                                     self.correct_obj,
-                                                                     self.correct_noobj,
-                                                                     self.tot_class_preds,
-                                                                     self.tot_obj,
-                                                                     self.tot_noobj)
-
-        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
-        self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True,on_step=False, on_epoch=True)
-        return loss
-
-
-    def evaluate(self, batch, stage=None):
-        x, y = batch
-        out = self(x)
-        loss = self.criterion(out,y,self.scaled_anchors)
-        acc  = self.accuracy_fn(y,out,self.hparams.threshold,self.correct_class,
-                                                                     self.correct_obj,
-                                                                     self.correct_noobj,
-                                                                     self.tot_class_preds,
-                                                                     self.tot_obj,
-                                                                     self.tot_noobj)
-
-        if stage:
-            self.log(f"{stage}_loss", loss, prog_bar=True)
-            self.log_dict({"class_accuracy": acc[0], "no_object_accuracy": acc[1], "object_accuracy":acc[2]},prog_bar=True)
-
-    def test_step(self, batch, batch_idx):
-        self.evaluate(batch, "test")
-
-    def validation_step(self, batch, batch_idx):
-        self.evaluate(batch, "val")
-
-    def configure_optimizers(self):
-        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr, weight_decay=self.hparams.weight_decay)
-        scheduler = OneCycleLR(
-                optimizer,
-                max_lr= 1E-3,
-                pct_start = 5/self.trainer.max_epochs,
-                epochs=self.trainer.max_epochs,
-                steps_per_epoch=len(train_loader),
-                div_factor=100,verbose=True,
-                three_phase=False
-            )
-        return ([optimizer],[scheduler])
-
-
+import litmodelclass
 
 
 # gradio
@@ -176,7 +71,7 @@ with gr.Blocks() as demo:
     def yolo3_inference(input_img,gradcam=True,gradcam_opa=0.5): # function for yolo inference
       
       # load model
-      yololit = LitYolo()
+      yololit = litmodelclass.LitYolo()
       inference_model = yololit.load_from_checkpoint("yolo3_model.ckpt")
       
       # bboxes, gradcam