Update app.py

app.py

@@ -19,6 +19,117 @@ label_names = [
     "Pulmonary_fibrosis"
     ]
 
+class VinDetector(pl.LightningModule):
+    def __init__(self, **kwargs):
+        super().__init__()
+
+        self.model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
+        num_classes = 15
+        in_features = self.model.roi_heads.box_predictor.cls_score.in_features
+        self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
+        
+        self.learning_rate = 1e-3
+        self.batch_size = 4
+
+    def forward(self, x):
+        return self.model(x)
+
+    def prepare_data(self):
+        
+        df = pd.read_csv('../input/vinbigdata-chest-xray-abnormalities-detection/train.csv')
+        df = df[df['class_id'] != 14].reset_index(drop=True)
+        self.train_dataset = VBDDataset(df, '../input/vinbigdata-chest-xray-original-png/train', get_train_transform())
+
+    def train_dataloader(self):
+        
+        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, pin_memory=True, num_workers=4, collate_fn=collate_fn)
+
+    def training_step(self, batch, batch_idx):
+        images, targets = batch
+        targets = [{k: v for k, v in t.items()} for t in targets]
+        
+        loss_dict = self.model(images, targets)
+        loss = sum(loss for loss in loss_dict.values())
+        self.log('Loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+        return {"loss": loss}
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.SGD(self.model.parameters(), lr=self.learning_rate, momentum=0.95, weight_decay=1e-5, nesterov=True)
+        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=6, eta_min=0, verbose=True)
+        return [optimizer], [scheduler]
+    
+    
+#VBDDataset Class
+class VBDDataset(Dataset):
+
+    def __init__(self, dataframe, image_dir, transforms=None, phase='train'):
+        super().__init__()
+
+        
+        self.image_ids = dataframe['image_id'].unique()
+        self.df = dataframe
+        self.image_dir = image_dir
+        self.transforms = transforms
+        self.phase = phase
+
+    def __getitem__(self, idx):
+        
+        image_id = self.image_ids[idx]
+        records = self.df[self.df['image_id'] == image_id]
+
+        
+        image = cv2.imread(f'{self.image_dir}/{image_id}.png', cv2.IMREAD_COLOR)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32)
+        image /= 255.0
+
+        
+        if self.phase == 'test':
+            if self.transforms:
+                sample = {
+                    'image': image,
+                }
+                sample = self.transforms(**sample)
+                image = sample['image']
+            return image, image_id
+
+        
+        boxes = records[['x_min', 'y_min', 'x_max', 'y_max']].values
+        
+    
+        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
+        area = torch.as_tensor(area, dtype=torch.float32)
+        
+    
+        labels = torch.squeeze(torch.as_tensor((records.class_id.values+1,), dtype=torch.int64))
+        iscrowd = torch.zeros((records.shape[0],), dtype=torch.int64)
+        
+        
+        target = {}
+        target['boxes'] = boxes
+        target['labels'] = labels
+        target['area'] = area
+        target['image_id'] = torch.tensor([idx])
+        target['iscrowd'] = iscrowd
+        
+        
+        if self.transforms:
+            sample = {
+                'image': image,
+                'bboxes': target['boxes'],
+                'labels': labels
+            }
+            sample = self.transforms(**sample)
+            image = sample['image']
+            
+            target['boxes'] = torch.as_tensor(sample['bboxes'])
+
+        return image, target
+
+    def __len__(self):
+
+        return self.image_ids.shape[0]
+    
+
 def generate_diagnostic_report(predictions, labels, threshold=0.5):
     # Initialize an empty report string
     report = "Diagnostic Report:\n\n"