import os
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
import torch
torch.rand(10)
import torch.nn as nn
import torch.nn.functional as F
import glob
from tqdm import tqdm, trange
print(torch.cuda.is_available())
print(torch.cuda.get_device_name())
print(torch.cuda.current_device())

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')
    
import torch.backends.cudnn as cudnn
import numpy as np
import os, cv2
from tqdm import tqdm, trange
import seaborn as sns
from models.experimental import attempt_load
from utils.datasets import LoadStreams, LoadImages
from utils.general import (
    check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer)
from utils.torch_utils import select_device, load_classifier, time_synchronized
from my_utils import xyxy_2_xyxyo, draw_boxes
# Initialize
device = select_device('')
half = device.type != 'cpu'  # half precision only supported on CUDA

def prepare_input(img1, img_size=416, half=True):
    
    img2 = cv2.resize(img1, (img_size, img_size)) # W x H
    img2 = img2.transpose(2,0,1)
    img2 = img2[np.newaxis, ...]
    img2 = torch.from_numpy(img2).to(device) # torch image is ch x H x W
    img2 = img2.half() if not half else img2.float()
    img2 /= 255.0
    
    return img2
#%%

# Directories
out = '/home/user01/data_ssd/Talha/yolo/op/'
weights = '/home/user01/data_ssd/Talha/yolo/ScaledYOLOv4/runs/exp2_yolov4-csp-results/weights/best_yolov4-csp-results.pt'
source = '/home/user01/data_ssd/Talha/yolo/paprika_y5/valid/images/'
imgsz = 416
conf_thres = 0.4
iou_thres = 0.5
classes =  [0,1,2,3,4,5] 
class_names = ["blossom_end_rot", "graymold","powdery_mildew","spider_mite",
               "spotting_disease", "snails_and_slugs"] 

# deleting files in op_dir
filelist = [ f for f in os.listdir(out)]# if f.endswith(".png") ]
for f in tqdm(filelist, desc = 'Deleting old files fro directory'):
    os.remove(os.path.join(out, f))

# Load model
model = attempt_load(weights, map_location=device)  # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size

if half:
    model.half()  # to FP16

# Load model
model = attempt_load(weights, map_location=device)  # load FP32 model
imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size

img_paths = glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png') + \
            glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')

# Run inference
if device.type != 'cpu':
    model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
#%%
for i in trange(len(img_paths)):
    
    
    path = img_paths[i]
    img1  = cv2.imread(path)
    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
    img_h, img_w, _ = img1.shape
    
    img2 = prepare_input(img1, 416, half)
    # get file name
    name = os.path.basename(path)[:-4]
    # Inference
    t1 = time_synchronized()
    pred = model(img2, augment=False)[0]

    # Apply NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=True)
    if pred[0] is not None:
        boxes = pred[0].cpu().detach().numpy() # <xmin><ymin><xmax><ymax><confd><class_id>
    else:
        boxes = np.array([10.0, 20.0, 30.0, 50.0, 0.75, 0]).reshape(1,6) # dummy values
    
    coords_minmax = np.zeros((boxes.shape[0], 4)) # droping 5th value
    confd = np.zeros((boxes.shape[0], 1))
    class_ids = np.zeros((boxes.shape[0], 1))
    # assign
    coords_minmax = boxes[:,0:4] # coords
    confd = boxes[:,4] # confidence
    class_ids = boxes[:,5] # class id
    
    coords_xyminmax = []
    det_classes = []
    for i in range(boxes.shape[0]):
        coords_xyminmax.append(xyxy_2_xyxyo(img_w, img_h, coords_minmax[i]))
        det_classes.append(class_names[int(class_ids[i])])
    all_bounding_boxnind = []
    for i in range(boxes.shape[0]):
        
        bounding_box = [0.0] * 6
        
        bounding_box[0] = det_classes[i]
        bounding_box[1] = confd[i]
        bounding_box[2] = coords_xyminmax[i][0]
        bounding_box[3] = coords_xyminmax[i][1]
        bounding_box[4] = coords_xyminmax[i][2]
        bounding_box[5] = coords_xyminmax[i][3]
        
        bounding_box = str(bounding_box)[1:-1]# remove square brackets
        bounding_box = bounding_box.replace("'",'')# removing inverted commas around class name
        bounding_box = "".join(bounding_box.split())# remove spaces in between **here dont give space inbetween the inverted commas "".
        all_bounding_boxnind.append(bounding_box)
    
    all_bounding_boxnind = ' '.join(map(str, all_bounding_boxnind))# convert list to string
    all_bounding_boxnind=list(all_bounding_boxnind.split(' ')) # convert strin to list
    # replacing commas with spaces
    for i in range(len(all_bounding_boxnind)):
        all_bounding_boxnind[i] = all_bounding_boxnind[i].replace(',',' ')
    for i in range(len(all_bounding_boxnind)):
    # check if file exiscts else make new
        with open(out +'{}.txt'.format(name), "a+") as file_object:
            # Move read cursor to the start of file.
            file_object.seek(0)
            # If file is not empty then append '\n'
            data = file_object.read(100)
            if len(data) > 0 :
                file_object.write("\n")
            # Append text at the end of file
            file_object.write(all_bounding_boxnind[i])
#%%
import glob, random
import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.dpi'] = 300

img_paths = glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png') + \
            glob.glob('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')
img_path = random.choice(img_paths)

img1  = cv2.imread(img_path)
img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
img_h, img_w, _ = img1.shape

img2 = prepare_input(img1, 416, half)

pred = model(img2, augment=False)[0]

# Apply NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=True)
boxes = pred[0].cpu().detach().numpy() # <xmin><ymin><xmax><ymax><confd><class_id>

coords_minmax = np.zeros((boxes.shape[0], 4)) # droping 5th value
confd = np.zeros((boxes.shape[0], 1))
class_ids = np.zeros((boxes.shape[0], 1))
# assign
coords_minmax = boxes[:,0:4] # coords
confd = boxes[:,4] # confidence
class_ids = boxes[:,5] # class id
    
coords_xyminmax = []
det_classes = []
for i in range(boxes.shape[0]):
    coords_xyminmax.append(xyxy_2_xyxyo(img_w, img_h, coords_minmax[i]))
    det_classes.append(class_names[int(class_ids[i])])

t = np.asarray(coords_xyminmax)
op = draw_boxes(img1, confd, t, det_classes, class_names, order='xy_minmax', analysis=False)
plt.imshow(op)
print('='*50)
print('Image Name: ', os.path.basename(img_path),img1.shape)
print('\nClass_name      ', '| B_box Coords ', '| Confidence')
print('_'*50)
for k in range(len(det_classes)):
    print(det_classes[k], t[k], confd[k])
print('='*50)
[((4, 0, 4, 14), 'torch.rand', 'torch.rand', ({(4, 11, 4, 13): '(10)'}, {}), '(10)', False, 'import torch\n'), ((36, 9, 36, 26), 'utils.torch_utils.select_device', 'select_device', ({(36, 23, 36, 25): '""""""'}, {}), "('')", False, 'from utils.torch_utils import select_device, load_classifier, time_synchronized\n'), ((64, 9, 64, 66), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm, trange\n'), ((68, 8, 68, 50), 'models.experimental.attempt_load', 'attempt_load', (), '', False, 'from models.experimental import attempt_load\n'), ((75, 8, 75, 50), 'models.experimental.attempt_load', 'attempt_load', (), '', False, 'from models.experimental import attempt_load\n'), ((161, 11, 161, 35), 'random.choice', 'random.choice', ({(161, 25, 161, 34): 'img_paths'}, {}), '(img_paths)', False, 'import glob, random\n'), ((163, 8, 163, 28), 'cv2.imread', 'cv2.imread', ({(163, 19, 163, 27): 'img_path'}, {}), '(img_path)', False, 'import os, cv2\n'), ((164, 7, 164, 44), 'cv2.cvtColor', 'cv2.cvtColor', ({(164, 20, 164, 24): 'img1', (164, 26, 164, 43): 'cv2.COLOR_BGR2RGB'}, {}), '(img1, cv2.COLOR_BGR2RGB)', False, 'import os, cv2\n'), ((172, 7, 172, 87), 'utils.general.non_max_suppression', 'non_max_suppression', (), '', False, 'from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer\n'), ((175, 16, 175, 45), 'numpy.zeros', 'np.zeros', ({(175, 25, 175, 44): '(boxes.shape[0], 4)'}, {}), '((boxes.shape[0], 4))', True, 'import numpy as np\n'), ((176, 8, 176, 37), 'numpy.zeros', 'np.zeros', ({(176, 17, 176, 36): '(boxes.shape[0], 1)'}, {}), '((boxes.shape[0], 1))', True, 'import numpy as np\n'), ((177, 12, 177, 41), 'numpy.zeros', 'np.zeros', ({(177, 21, 177, 40): '(boxes.shape[0], 1)'}, {}), '((boxes.shape[0], 1))', True, 'import numpy as np\n'), ((189, 4, 189, 31), 'numpy.asarray', 'np.asarray', ({(189, 15, 189, 30): 'coords_xyminmax'}, {}), '(coords_xyminmax)', True, 'import numpy as np\n'), ((190, 5, 190, 92), 'my_utils.draw_boxes', 'draw_boxes', (), '', False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((191, 0, 191, 14), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(191, 11, 191, 13): 'op'}, {}), '(op)', True, 'import matplotlib.pyplot as plt\n'), ((9, 6, 9, 31), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((10, 6, 10, 34), 'torch.cuda.get_device_name', 'torch.cuda.get_device_name', ({}, {}), '()', False, 'import torch\n'), ((11, 6, 11, 33), 'torch.cuda.current_device', 'torch.cuda.current_device', ({}, {}), '()', False, 'import torch\n'), ((41, 11, 41, 49), 'cv2.resize', 'cv2.resize', ({(41, 22, 41, 26): 'img1', (41, 28, 41, 48): '(img_size, img_size)'}, {}), '(img1, (img_size, img_size))', False, 'import os, cv2\n'), ((78, 12, 78, 86), 'glob.glob', 'glob.glob', ({(78, 22, 78, 85): '"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png"""'}, {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png')", False, 'import glob, random\n'), ((79, 12, 79, 86), 'glob.glob', 'glob.glob', ({(79, 22, 79, 85): '"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg"""'}, {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')", False, 'import glob, random\n'), ((89, 12, 89, 28), 'cv2.imread', 'cv2.imread', ({(89, 23, 89, 27): 'path'}, {}), '(path)', False, 'import os, cv2\n'), ((90, 11, 90, 48), 'cv2.cvtColor', 'cv2.cvtColor', ({(90, 24, 90, 28): 'img1', (90, 30, 90, 47): 'cv2.COLOR_BGR2RGB'}, {}), '(img1, cv2.COLOR_BGR2RGB)', False, 'import os, cv2\n'), ((97, 9, 97, 28), 'utils.torch_utils.time_synchronized', 'time_synchronized', ({}, {}), '()', False, 'from utils.torch_utils import select_device, load_classifier, time_synchronized\n'), ((101, 11, 101, 91), 'utils.general.non_max_suppression', 'non_max_suppression', (), '', False, 'from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords, xyxy2xywh, plot_one_box, strip_optimizer\n'), ((107, 20, 107, 49), 'numpy.zeros', 'np.zeros', ({(107, 29, 107, 48): '(boxes.shape[0], 4)'}, {}), '((boxes.shape[0], 4))', True, 'import numpy as np\n'), ((108, 12, 108, 41), 'numpy.zeros', 'np.zeros', ({(108, 21, 108, 40): '(boxes.shape[0], 1)'}, {}), '((boxes.shape[0], 1))', True, 'import numpy as np\n'), ((109, 16, 109, 45), 'numpy.zeros', 'np.zeros', ({(109, 25, 109, 44): '(boxes.shape[0], 1)'}, {}), '((boxes.shape[0], 1))', True, 'import numpy as np\n'), ((159, 12, 159, 86), 'glob.glob', 'glob.glob', ({(159, 22, 159, 85): '"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png"""'}, {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.png')", False, 'import glob, random\n'), ((160, 12, 160, 86), 'glob.glob', 'glob.glob', ({(160, 22, 160, 85): '"""/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg"""'}, {}), "('/home/user01/data_ssd/Talha/yolo/paprika_y5/test/images/*.jpg')", False, 'import glob, random\n'), ((193, 22, 193, 48), 'os.path.basename', 'os.path.basename', ({(193, 39, 193, 47): 'img_path'}, {}), '(img_path)', False, 'import os, cv2\n'), ((13, 32, 13, 57), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((19, 10, 19, 39), 'torch.cuda.get_device_name', 'torch.cuda.get_device_name', ({(19, 37, 19, 38): '(0)'}, {}), '(0)', False, 'import torch\n'), ((63, 24, 63, 39), 'os.listdir', 'os.listdir', ({(63, 35, 63, 38): 'out'}, {}), '(out)', False, 'import os, cv2\n'), ((65, 14, 65, 34), 'os.path.join', 'os.path.join', ({(65, 27, 65, 30): 'out', (65, 32, 65, 33): 'f'}, {}), '(out, f)', False, 'import os, cv2\n'), ((95, 11, 95, 33), 'os.path.basename', 'os.path.basename', ({(95, 28, 95, 32): 'path'}, {}), '(path)', False, 'import os, cv2\n'), ((186, 27, 186, 71), 'my_utils.xyxy_2_xyxyo', 'xyxy_2_xyxyo', ({(186, 40, 186, 45): 'img_w', (186, 47, 186, 52): 'img_h', (186, 54, 186, 70): 'coords_minmax[i]'}, {}), '(img_w, img_h, coords_minmax[i])', False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((44, 11, 44, 33), 'torch.from_numpy', 'torch.from_numpy', ({(44, 28, 44, 32): 'img2'}, {}), '(img2)', False, 'import torch\n'), ((118, 31, 118, 75), 'my_utils.xyxy_2_xyxyo', 'xyxy_2_xyxyo', ({(118, 44, 118, 49): 'img_w', (118, 51, 118, 56): 'img_h', (118, 58, 118, 74): 'coords_minmax[i]'}, {}), '(img_w, img_h, coords_minmax[i])', False, 'from my_utils import xyxy_2_xyxyo, draw_boxes\n'), ((21, 30, 21, 60), 'torch.cuda.memory_allocated', 'torch.cuda.memory_allocated', ({(21, 58, 21, 59): '(0)'}, {}), '(0)', False, 'import torch\n'), ((22, 30, 22, 59), 'torch.cuda.memory_reserved', 'torch.cuda.memory_reserved', ({(22, 57, 22, 58): '(0)'}, {}), '(0)', False, 'import torch\n'), ((105, 16, 105, 59), 'numpy.array', 'np.array', ({(105, 25, 105, 58): '[10.0, 20.0, 30.0, 50.0, 0.75, 0]'}, {}), '([10.0, 20.0, 30.0, 50.0, 0.75, 0])', True, 'import numpy as np\n'), ((83, 10, 83, 41), 'torch.zeros', 'torch.zeros', ({(83, 22, 83, 23): '(1)', (83, 25, 83, 26): '(3)', (83, 28, 83, 33): 'imgsz', (83, 35, 83, 40): 'imgsz'}, {}), '(1, 3, imgsz, imgsz)', False, 'import torch\n')]