Update app.py

app.py

@@ -60,77 +60,6 @@ class VinDetector(pl.LightningModule):
         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 get_train_transform():
     return A.Compose([