【Data Enhancement】Brightness/Contrast Enhancement

1. Effect diagram

insert image description here

2. Code

import cv2
import math
import numpy as np
import os
import glob
import json
import shutil
import xml.etree.ElementTree as ET
from xml.etree.ElementTree import ElementTree, Element

def getColorImg(alpha,beta,img_path,img_write_path):
    img = cv2.imread(img_path)
    colored_img = np.uint8(np.clip((alpha * img + beta), 0, 255))
    cv2.imwrite(img_write_path,colored_img)

def getColorAnno(anno_path,anno_write_path):
    tree = ET.parse(anno_path)
    tree.write(anno_write_path)  # 保存修改后的XML文件

def color(alpha,beta,img_dir,anno_dir,img_write_dir,anno_write_dir):
    if not os.path.exists(img_write_dir):
        os.makedirs(img_write_dir)

    if not os.path.exists(anno_write_dir):
        os.makedirs(anno_write_dir)
    img_names=os.listdir(img_dir)
    for img_name in img_names:
        img_path=os.path.join(img_dir,img_name)
        img_write_path=os.path.join(img_write_dir,img_name[:-4]+'color'+str(int(alpha*10))+'.jpg')
        #
        anno_path=os.path.join(anno_dir,img_name[:-4]+'.xml')
        anno_write_path = os.path.join(anno_write_dir, img_name[:-4]+'color'+str(int(alpha*10))+'.xml')
        #
        getColorImg(alpha,beta,img_path,img_write_path)
        getColorAnno(anno_path,anno_write_path)

alphas = [0.3, 0.5, 1.2]
beta = 10
# 原始图片地址
img_dir = r'J:\hive-master\datasets\1\images'
# 原始标签地址
anno_dir = r'J:\hive-master\datasets\1\annotations'
# 对比增强后图片地址
img_write_dir = r'J:\hive-master\datasets\1\color\color_imgs'
# 对比增强后标签地址
anno_write_dir = r'J:\hive-master\datasets\1\color\color_annotations'
for alpha in alphas:
    color(alpha, beta, img_dir, anno_dir, img_write_dir, anno_write_dir)

Guess you like

Origin blog.csdn.net/qq_38253797/article/details/123021160
Recommended