Update README.md

detection_models/yolo_stamp/train.ipynb

@@ -1,185 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from model import *\n",
-    "from loss import *\n",
-    "from data import *\n",
-    "from torch import optim\n",
-    "from tqdm import tqdm\n",
-    "\n",
-    "import pytorch_lightning as pl\n",
-    "from torchmetrics.detection import MeanAveragePrecision\n",
-    "from pytorch_lightning.loggers import TensorBoardLogger"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "_, _, test_dataset = get_datasets()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class LitModel(pl.LightningModule):\n",
-    "    def __init__(self):\n",
-    "        super().__init__()\n",
-    "        self.model = YOLOStamp()\n",
-    "        self.criterion = YOLOLoss()\n",
-    "        self.val_map = MeanAveragePrecision(box_format='xywh', iou_type='bbox')\n",
-    "    \n",
-    "    def forward(self, x):\n",
-    "        return self.model(x)\n",
-    "\n",
-    "    def configure_optimizers(self):\n",
-    "        optimizer = optim.AdamW(self.parameters(), lr=1e-3)\n",
-    "        # return optimizer\n",
-    "        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, 1000)\n",
-    "        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler}\n",
-    "\n",
-    "    def training_step(self, batch, batch_idx):\n",
-    "        images, targets = batch\n",
-    "        tensor_images = torch.stack(images)\n",
-    "        tensor_targets = torch.stack(targets)\n",
-    "        output = self.model(tensor_images)\n",
-    "        loss = self.criterion(output, tensor_targets)\n",
-    "        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n",
-    "        return loss\n",
-    "\n",
-    "    def validation_step(self, batch, batch_idx):\n",
-    "        images, targets = batch\n",
-    "        tensor_images = torch.stack(images)\n",
-    "        tensor_targets = torch.stack(targets)\n",
-    "        output = self.model(tensor_images)\n",
-    "        loss = self.criterion(output, tensor_targets)\n",
-    "        self.log(\"val_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n",
-    "\n",
-    "        for i in range(len(images)):\n",
-    "            boxes = output_tensor_to_boxes(output[i].detach().cpu())\n",
-    "            boxes = nonmax_suppression(boxes)\n",
-    "            target = target_tensor_to_boxes(targets[i])[::BOX]\n",
-    "            if not boxes:\n",
-    "                boxes = torch.zeros((1, 5))\n",
-    "            preds = [\n",
-    "                dict(\n",
-    "                    boxes=torch.tensor(boxes)[:, :4].clone().detach(),\n",
-    "                    scores=torch.tensor(boxes)[:, 4].clone().detach(),\n",
-    "                    labels=torch.zeros(len(boxes)),\n",
-    "                )\n",
-    "            ]\n",
-    "            target = [\n",
-    "                dict(\n",
-    "                    boxes=torch.tensor(target),\n",
-    "                    labels=torch.zeros(len(target)),\n",
-    "                )\n",
-    "            ]\n",
-    "            self.val_map.update(preds, target)\n",
-    "    \n",
-    "    def on_validation_epoch_end(self):\n",
-    "        mAPs = {\"val_\" + k: v for k, v in self.val_map.compute().items()}\n",
-    "        mAPs_per_class = mAPs.pop(\"val_map_per_class\")\n",
-    "        mARs_per_class = mAPs.pop(\"val_mar_100_per_class\")\n",
-    "        self.log_dict(mAPs)\n",
-    "        self.val_map.reset()\n",
-    "\n",
-    "        image = test_dataset[randint(0, len(test_dataset) - 1)][0].to(self.device)\n",
-    "        output = self.model(image.unsqueeze(0))\n",
-    "        boxes = output_tensor_to_boxes(output[0].detach().cpu())\n",
-    "        boxes = nonmax_suppression(boxes)\n",
-    "        img = image.permute(1, 2, 0).cpu().numpy()\n",
-    "        img = visualize_bbox(img.copy(), boxes=boxes)\n",
-    "        img = (255. * (img * np.array(STD) + np.array(MEAN))).astype(np.uint8)\n",
-    "        \n",
-    "        self.logger.experiment.add_image(\"detected boxes\", torch.tensor(img).permute(2, 0, 1), self.current_epoch)\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "litmodel = LitModel()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "logger = TensorBoardLogger(\"detection_logs\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "epochs = 100"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_loader, val_loader = get_loaders(batch_size=8)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "trainer = pl.Trainer(accelerator=\"auto\", max_epochs=epochs, logger=logger)\n",
-    "trainer.fit(model=litmodel, train_dataloaders=train_loader, val_dataloaders=val_loader)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%tensorboard"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.9.0"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

detection_models/yolo_stamp/utils.py

@@ -1,5 +1,4 @@
 import torch
-import cv2
 import pandas as pd
 import numpy as np
 from pathlib import Path
@@ -53,33 +52,6 @@ def plot_normalized_img(img, std=STD, mean=MEAN, size=(7,7)):
     plt.figure(figsize=size)
     plt.imshow((255. * (img * std + mean)).astype(np.uint))
     plt.show()
-    
-
-def visualize_bbox(img, boxes, thickness=2, color=BOX_COLOR, draw_center=True):
-    """
-        Draws boxes on the given image.
-
-        Arguments:
-        img -- torch.Tensor of shape (3, W, H) or numpy.ndarray of shape (W, H, 3)
-        boxes -- list of shape (None, 5)
-        thickness -- number specifying the thickness of box border
-        color -- RGB tuple of shape (3,) specifying the color of boxes
-        draw_center -- boolean specifying whether to draw center or not
-
-        Returns:
-        img_copy -- numpy.ndarray of shape(W, H, 3) containing image with bouning boxes
-    """
-    img_copy = img.cpu().permute(1,2,0).numpy() if isinstance(img, torch.Tensor) else img.copy()
-    for box in boxes:
-        x,y,w,h = int(box[0]), int(box[1]), int(box[2]), int(box[3])
-        img_copy = cv2.rectangle(
-            img_copy,
-            (x,y),(x+w, y+h),
-            color, thickness)
-        if draw_center:
-            center = (x+w//2, y+h//2)
-            img_copy = cv2.circle(img_copy, center=center, radius=3, color=(0,255,0), thickness=2)
-    return img_copy
 
 
 def read_data(annotations=Path(ANNOTATIONS_PATH)):

requirements.txt

@@ -1,7 +1,14 @@
-torch==1.12.0
-torchvision==0.13.0
-scikit-learn==1.1.3
+albumentations==1.3.0
+click==8.0.4
+gradio==3.36.1
+huggingface_hub==0.14.1
 matplotlib==3.6.0
-pillow==9.3.0
+numpy==1.23.4
 pandas==1.5.1
-gradio>=3.36.1
+Pillow==9.3.0
+Pillow==10.0.0
+pytorch_lightning==2.0.2
+scikit_learn==1.1.3
+torch==1.12.0+cu116
+torchvision==0.13.0+cu116
+tqdm==4.64.1