draw_box.py

import cv2
from matplotlib.pyplot import imshow
import pandas as pd
import matplotlib.pyplot as plt

class draw_object_detection:

  def __init__(self,file_name):
    self.file_name = file_name 
    self.file_data = self.read_file()
    self.img_path = []
    self.img_data = []
    self.img_label = []
    self.xy_min = []
    self.xy_max = []

  def read_file(self):
    df = pd.read_csv(self.file_name)
    return df.to_numpy()

  def info(self,index=0):
    base_path = "data/test/"
    temp = self.file_data[index]
    self.img_path = base_path + temp[0]
    self.img_label = temp[1]
    self.img_score = temp[2]
    self.img_box = temp[3:7]
    self.xy_min = tuple(temp[3:5])
    self.xy_max = tuple(temp[5:7])

    # print(temp)
    print("path: ",self.img_path)
    print("label: ",self.img_label)
    print("score: ",self.img_score)

    # xmin	ymin	xmax	ymax
    print("box: ",self.img_box)
    print("xy_min: ",self.xy_min)
    print("xy_max: ",self.xy_max)

    self.img_data = self.read()

  def view(self):
    return imshow(self.img_data)

  def read(self):
    image = cv2.imread(self.img_path)
    return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

  def draw(self):
    t1 = cv2.rectangle(self.img_data, self.xy_min, self.xy_max,color=(0, 255, 0), thickness=3)
    t2 = cv2.cvtColor(t1, cv2.COLOR_BGR2RGB)
    self.img_data = t2

  # view img before and after
  def pp(self):
    raw = self.read()
    self.draw()
    changed = self.img_data

    font_size = 15

    for num in range(1):
      plt.figure()
      plt.figure(figsize = (18,40))
      plt.subplot(121)
      plt.title("Img: "+self.img_path,fontdict = {'fontsize' : font_size})
      plt.imshow(raw)
      plt.subplot(122)
      plt.title("Detected: "+self.img_label+" | Score: "+str(self.img_score),fontdict = {'fontsize' : font_size})
      plt.imshow(changed)
      plt.show()


a = draw_object_detection("submission.csv")
a.info(9)
# a.draw()
# a.view()
a.pp()


# submission file format
# image_id,class,score,xmin,ymin,xmax,ymax,width,height