import os
import cv2
import numpy as np
import torch
import torchvision.models as models

SPORTS_BALL_CLASS = 37

model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval().cuda()
image_folder = 'small_images/ball/Ball2-small/images'
output_folder = ''


def generate_label(image_name):
    # read image
    print(f'reading image: {image_name}')
    raw_img = cv2.imread(os.path.join(image_folder, image_name))
    cv2.imshow('test', raw_img)
    cv2.waitKey(0)
    # process it
    img = np.moveaxis(raw_img, -1, 0)
    img = img / 255
    img = torch.tensor([img]).float().cuda()
    
    output = model(img)

    boxes = output[0]['boxes']
    labels = output[0]['labels']
    scores = output[0]['scores']
    
    box = None
    for box, label, score in zip(boxes, labels, scores):
        if label == SPORTS_BALL_CLASS:
            box, score = box.tolist(), score.tolist()
            box = list(map(int, box))
            cv2.rectangle(raw_img, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2)
            cv2.imshow('test', raw_img)
            cv2.waitKey(0)
            break
    else:
        print(image_name, 'cant see ball')
        box = None # haven't found anything


    # [format: "label::annotation_type|filename|img_width|img_height|x1|y1|x2|y2|center_x|center_y|width|height"]
    # label::ball|img_fake_cam_008792.PNG|1920|1080|890|162|966|242|928.0|202.0|76|80
    if box is not None:
        return 'label::ball|{file_name}|0|0|{x1}|{y1}|{x2}|{y2}|0|0|0|0\n'.format(
            file_name=image_name,
            x1=int(box[0]),
            y1=int(box[1]),
            x2=int(box[2]),
            y2=int(box[3]))
    return ''

def main():
    with open('annotations.txt', 'w') as annotations:
        annotations.write('This is autogenerated file\n\n\n\n\n\n')
        
        for file in os.listdir(image_folder):
            label = generate_label(file)
            annotations.write(label)
        

if __name__ == '__main__':
    main()

import os
import cv2
import numpy as np
import torch
import torchvision.models as models

SPORTS_BALL_CLASS = 37

model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model.eval().cuda()
image_folder = 'small_images/ball/Ball2-small/images'
output_folder = ''


def generate_label(image_name):
    # read image
    print(f'reading image: {image_name}')
    raw_img = cv2.imread(os.path.join(image_folder, image_name))
    cv2.imshow('test', raw_img)
    cv2.waitKey(0)
    # process it
    img = np.moveaxis(raw_img, -1, 0)
    img = img / 255
    img = torch.tensor([img]).float().cuda()
    
    output = model(img)

    boxes = output[0]['boxes']
    labels = output[0]['labels']
    scores = output[0]['scores']
    
    box = None
    for box, label, score in zip(boxes, labels, scores):
        if label == SPORTS_BALL_CLASS:
            box, score = box.tolist(), score.tolist()
            box = list(map(int, box))
            cv2.rectangle(raw_img, (box[0], box[1]), (box[2], box[3]), (255,0,0), 2)
            cv2.imshow('test', raw_img)
            cv2.waitKey(0)
            break
    else:
        print(image_name, 'cant see ball')
        box = None # haven't found anything


    # [format: "label::annotation_type|filename|img_width|img_height|x1|y1|x2|y2|center_x|center_y|width|height"]
    # label::ball|img_fake_cam_008792.PNG|1920|1080|890|162|966|242|928.0|202.0|76|80
    if box is not None:
        return 'label::ball|{file_name}|0|0|{x1}|{y1}|{x2}|{y2}|0|0|0|0\n'.format(
            file_name=image_name,
            x1=int(box[0]),
            y1=int(box[1]),
            x2=int(box[2]),
            y2=int(box[3]))
    return ''

def main():
    with open('annotations.txt', 'w') as annotations:
        annotations.write('This is autogenerated file\n\n\n\n\n\n')
        
        for file in os.listdir(image_folder):
            label = generate_label(file)
            annotations.write(label)
        

if __name__ == '__main__':
    main()

# read image # process it # haven't found anything # [format: "label::annotation_type|filename|img_width|img_height|x1|y1|x2|y2|center_x|center_y|width|height"] # label::ball|img_fake_cam_008792.PNG|1920|1080|890|162|966|242|928.0|202.0|76|80