Update app.py

app.py

@@ -6,6 +6,8 @@ from torchvision.utils import draw_bounding_boxes
 import numpy as np
 from PIL import Image, ImageDraw
 import torchvision.transforms.functional as F
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
 
 # Định nghĩa các lớp và màu sắc
 CLASS_NAMES = ['bone-fracture', 'angle', 'fracture', 'line', 'messed_up_angle']
@@ -59,6 +61,24 @@ def predict_image(image, model, device, transform, threshold=0.5):
 
     return drawn_image.permute(1, 2, 0)  # Đảm bảo định dạng hình ảnh là (H, W, C)
 
+# Định nghĩa hàm get_transforms
+def get_transforms(train=False):
+    if train:
+        transform = A.Compose([
+            A.Resize(800, 800),  # Kích thước đầu vào 800px
+            A.HorizontalFlip(p=0.3),
+            A.VerticalFlip(p=0.3),
+            A.RandomBrightnessContrast(p=0.1),
+            A.ColorJitter(p=0.1),
+            ToTensorV2()
+        ], bbox_params=A.BboxParams(format='coco'))
+    else:
+        transform = A.Compose([
+            A.Resize(800, 800),  # Kích thước đầu vào 800px
+            ToTensorV2()
+        ], bbox_params=A.BboxParams(format='coco'))
+    return transform
+
 # Tải mô hình
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = fasterrcnn_resnet50_fpn(pretrained=False)  # Khởi tạo mô hình
@@ -90,4 +110,4 @@ iface = gr.Interface(
 
 # Chạy ứng dụng
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(share=True)  # Set share=True to create a public link